import numpy as np
import torch
from transformers import BertTokenizer, BertModel

# 加载预训练的BERT模型和tokenizer
model_name = 'E:/phythonproject/课设/bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertModel.from_pretrained(model_name)

# 预处理文本并生成BERT输入
def preprocess_text(text):
    encoded_input = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=128,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )
    return encoded_input

# 使用BERT模型获取文本的表示向量
def get_text_embedding(text):
    inputs = preprocess_text(text)#使用preprocess_text(text)函数的值作为输入
    with torch.no_grad():
        outputs = model(**inputs)
        embedding = torch.mean(outputs.last_hidden_state, dim=1)  # 使用平均池化获取 句子向量
    return embedding.numpy().flatten()  # 转换为numpy数组,并展平为1维向量

# 计算两个文本之间的相似度
def calculate_similarity(text1, text2):
    embedding1 = get_text_embedding(text1) #获取文本的表示向量
    embedding2 = get_text_embedding(text2) #获取文本的表示向量
    similarity = np.dot(embedding1, embedding2) / (np.linalg.norm(embedding1) * np.linalg.norm(embedding2)) #计算他们的点积再除以范数
    return similarity   #返回相似度值

# 示例使用
text1 = '小米8 全面屏游戏智能手机 6GB+128GB 黑色 全网通4G 双卡双待 拍照手机'
texts = [
    '【必购码】小米8 全面屏游戏智能手机 6GB+128GB 黑色 全网通4G 双卡双待',
    '【必购码】小米8 全面屏游戏智能手机 6GB+64GB 黑色 全网通4G 双卡双待',
    '小米MIX2 全面屏游戏手机 6GB+128GB 黑色 全网通4G手机 双卡双待',
    '小米8 全面屏游戏智能手机 6GB+64GB 蓝色 全网通4G 双卡双待',
    '小米8SE 全面屏智能手机 6GB+64GB 金色 全网通4G 双卡双待 拍照手机 游戏手机',
    '【必购码】小米8 全面屏游戏智能手机 6GB+64GB 白色 全网通4G 双卡双待',
    '【必购码】小米8 全面屏游戏智能手机 6GB+128GB 白色 全网通4G 双卡双待',
    '【必购码】小米8 全面屏游戏智能手机 6GB+256GB 黑色 全网通4G 双卡双待',
    '【电信赠费版】小米8 全面屏游戏智能手机 6GB 64GB 白色 全网通4G 双卡双待'
]

print(text1 + '  1.0')
for text in texts:
    similarity = calculate_similarity(text1, text)
    print(text + '  ' + str(similarity))
